import matplotlib.pyplot as plt
import numpy as np
import paddle
from paddle.nn import Conv2D
from paddle.nn.initializer import Assign

class ImageProcessor:
    def __init__(self):
        # 设置matplotlib支持中文显示
        plt.rcParams['font.sans-serif'] = ['SimHei']  # 指定默认字体为黑体
        plt.rcParams['axes.unicode_minus'] = False  # 解决保存图像是负号'-'显示为方块的问题

    def generate_bw_image(self, black_ratio=0.6):
        """
        生成黑白图片并进行卷积操作，允许用户调整黑白比例。
        
        参数:
        black_ratio (float): 黑色区域的比例，默认为0.6（即60%的区域为黑色）。
        """
        # 确保 black_ratio 在 0 到 1 之间
        if not 0 <= black_ratio <= 1:
            raise ValueError("black_ratio 必须在 0 到 1 之间")

        # 创建初始化权重参数w
        w = np.array([1, 0, -1], dtype='float32')
        # 将权重参数调整成维度为[cout, cin, kh, kw]的四维张量
        w = w.reshape([1, 1, 1, 3])
        # 创建卷积算子，设置输出通道数，卷积核大小，和初始化权重参数
        conv = Conv2D(in_channels=1, out_channels=1, kernel_size=[1, 3],
            weight_attr=paddle.ParamAttr(
                initializer=Assign(value=w)))

        # 创建输入图片，图片左边的像素点取值为1，右边的像素点取值为0
        img = np.ones([50, 50], dtype='float32')
        # 根据用户传入的黑白比例调整图片
        split_idx = int(50 * black_ratio)
        img[:, split_idx:] = 0.

        # 将图片形状调整为[N, C, H, W]的形式
        x = img.reshape([1, 1, 50, 50])
        # 将numpy.ndarray转化成paddle中的tensor
        x = paddle.to_tensor(x)
        # 使用卷积算子作用在输入图片上
        y = conv(x)
        # 将输出tensor转化为numpy.ndarray
        out = y.numpy()

        # 显示输入图片和输出特征图
        f = plt.subplot(121)
        f.set_title('输入图像', fontsize=15)
        plt.imshow(img, cmap='gray')
        f = plt.subplot(122)
        f.set_title('输出特征图', fontsize=15)
        plt.imshow(out.squeeze(), cmap='gray')
        plt.show()

